package optimization;

import static optimization.LinearAlgebra.*;
import theory.vanpelt.bifurcation.PDistEstimate;

import common.statistics.Series;

public class SPSA {
	
	private static double[] bernoulli(int n, double[] w)
	{
		double[] r = new double[n];
		for (int i = 0; i < n; i++) r[i] = (Math.random() < 0.50 ? +1 : -1) * w[i];
		return r;
	}
	
	public static double[] SPSASolver(VectorFunction loss, double[] x0, double[] y)
	{
		// STEP 0 : Initialization
		int k = 0; // set index k = 0
		// pick initial guess x0 (theta 0)
		// select a,c,A,alpha,gamma
		double alpha = 0.602, gamma = 0.101; // From Spall2003
		double a=0, c=0.2, A=10;
		double[] x = x0;
		double maxDX = 1.0;
		
		//  0.1 = change in x0, 10 = magn of g0
		a = (0.1/10)*Math.pow(1+A, alpha);
//		a = 0.1;
		
		int maxiter=100;
		
		while(maxiter-- > 0) {
			
			double ak = a / Math.pow(k+1+A,alpha);
			double ck = c / Math.pow(k+1, gamma);
			
			double[] gradSum = new double[x0.length];
			for (int i = 0; i < 4; i++) {
				// STEP 1 : Generate perturbation vector
				double[] dk = bernoulli(x0.length, new double[]{ 0.001, 0.001 });
				double[] xplus  = vAdd(x, vMul(ck, dk));
				double[] xminus = vSub(x, vMul(ck, dk));
				
				// STEP 2 : Obtain loss function measurements
				double[] ypk = loss.evaluate(xplus);
				double[] ymk = loss.evaluate(xminus);
	
				// STEP 3: Gradient approximation
				double[] gradApprx = vOuterMul(vSub(ypk, ymk), vRecip(vMul(2*ck, dk)));
				gradSum = vAdd(gradSum, gradApprx);
				// *** do average of multiple approximations if noise is high
			}
			double[] gradApprx = vMul(1/10.0, gradSum);
			
			// STEP 4: Update x est
			double[] nx = vSub(x, vMul(ak, gradApprx)); // nx = x + a*g;

			// STEP 4.b: check constraints
			if (nx[0] < 0) nx[0] = 0;
			if (nx[0] > 6) nx[0] = 6;
			if (nx[1] < 0) nx[1] = 0;
			if (nx[1] > 1) nx[1] = 1;
			
			double[] cyloss = loss.evaluate(x);
			double[] nyloss = loss.evaluate(nx);
//			System.out.println("CYLOSS " + cyloss + ", NYLOSS " + nyloss);
			if (vNorm2(vSub(x, nx)) < maxDX) {
				if (vNorm2(nyloss) < vNorm2(cyloss)*1.2) x = nx;
			}
//			x = nx;
			
			
			System.out.println(k + " " + vPrint(x) + " " + vPrint(nyloss) + " " + ak);
			// STEP 5: Iteration or termination
			// terminate if little change or max iters
			if (vNorm2(nyloss) < 0.001) break; // no change
			k++;
		}
		
		return x;
	}

	
	public static double[] adaptiveSPSASolver(VectorFunction f, double[] x0, double[] y)
	{
		// STEP 0 : Initialization
		int k = 0; // set index k = 0
		// pick initial guess x0 (theta 0)
		// select a,c,A,alpha,gamma
		double alpha = 0.602, gamma = 0.101; // From Spall2003
		double a=0, c=0.5, ct=0.5, A=10;
		double[] x = x0;
		double maxDX = 1.0;
		double[] H = new double[x0.length*x0.length];
		
		//  0.1 = change in x0, 10 = magn of g0
//		a = (0.1/10)*Math.pow(1+A, alpha);
		a = 0.1;
		
		int maxiter=100;
		
		while(maxiter-- > 0) {
			
			double ak = a / (k+1);
			double ck = c / Math.pow(k+1, gamma);
			double ctk = ct / Math.pow(k+1, gamma);
			
			// STEP 1 : Generate perturbation vector
			double[] dk = bernoulli(x0.length, new double[]{ 0.001, 0.001 });
			double[] xplus  = vAdd(x, vMul(ck, dk));
			double[] xminus = vSub(x, vMul(ck, dk));
			
			// STEP 2 : Gradient approx 
			double[] ypk = f.evaluate(xplus);
			double[] ymk = f.evaluate(xminus);
			double[] gradApprx = vOuterMul(vSub(ypk, ymk), vRecip(vMul(2*ck, dk)));
			
			// *** do average of multiple approximations if noise is high

			// STEP 3: Hessian approx update
			double[] dtk = bernoulli(x0.length, new double[]{ 0.001, 0.001 });
			double[] xplus2  = vAdd(xplus, vMul(ctk, dtk));
			double[] Gkp = vOuterMul(vMul(1/ctk, vSub(f.evaluate(xplus2), ypk)), vRecip(dtk));
			double[] xminus2  = vSub(xminus, vMul(ctk, dtk));
			double[] Gkm = vOuterMul(vMul(1/ctk, vSub(f.evaluate(xminus2), ymk)), vRecip(dtk));

			double[] dGk = vSub(Gkp, Gkm); // G(x+cd) - G(x-cd)
			double[] Gkcdk = vMul(1/(2*ctk), vOuterMul(dGk, vRecip(dtk)));
			double[] Ht = vMul(0.5, vAdd(Gkcdk, mTrans(Gkcdk))); // 1/2 * (G + G')
			
			double kk = k / (k+1);
			double k1 = 1 / (k+1);
			double[] Hk = vAdd(vMul(kk, H), vMul(k1, Ht));
			
			// STEP 4: Update x est
			double[] nx = vSub(x, vMul(ak, mvMul(mInv(Hk), gradApprx))); // nx = x + a*H^-1*g;

			// STEP 4.b: check constraints
			if (nx[0] < 0) nx[0] = 0;
			if (nx[0] > 6) nx[0] = 6;
			if (nx[1] < 0) nx[1] = 0;
			if (nx[1] > 1) nx[1] = 1;
			
//			double[] cyloss = loss.evaluate(x);
//			double[] nyloss = loss.evaluate(nx);
//			System.out.println("CYLOSS " + cyloss + ", NYLOSS " + nyloss);
//			if (vNorm2(vSub(x, nx)) < maxDX) {
//				if (vNorm2(nyloss) < vNorm2(cyloss)*1.2) x = nx;
//			}
			x = nx;
			
			
			System.out.println(k + " " + vPrint(x));// + " " + vPrint(nyloss) + " " + ak);
			// STEP 5: Iteration or termination
			// terminate if little change or max iters
//			if (vNorm2(nyloss) < 0.001) break; // no change
			k++;
		}
		
		return null;
	}

	
	public static void main(String[] args) {
		
		
		VectorFunction f = new VectorFunction() {
			@Override
			public int yDim() {
				return 2;
			}
			@Override
			public double[] evaluate(double[] p) {
				PDistEstimate pd = new PDistEstimate(p[0], p[1]);
				// exact
				return new double[]{ pd.getMeanEst(), pd.getDevEst() };
//				Series s = pd.sample(14000);
	//			return new double[]{ s.mean(), s.stdevp() };
			}
		};
		
		VectorFunction loss = new VectorFunction() {
			@Override
			public int yDim() {
				return 1;
			}
			@Override
			public double[] evaluate(double[] p) {
				PDistEstimate pd = new PDistEstimate(p[0], p[1]);
				// exact
				return new double[] { Math.sqrt( Math.pow(pd.getMeanEst()-5.35,2) + Math.pow(pd.getDevEst()-4.41,2)) };
//				return new double[]{ pd.getMeanEst(), pd.getDevEst() };
//				Series s = pd.sample(14000);
	//			return new double[]{ s.mean(), s.stdevp() };
			}
		};
		SPSA.SPSASolver(loss, new double[]{2.5, 0.5}, new double[]{5.35, 4.41});
//		SPSA.adaptiveSPSASolver(loss, new double[]{2.5, 0.5}, new double[]{5.35, 4.41});
	}
}
